#!/usr/bin/env python3
"""
PyG 算子性能测试运行脚本

快速运行性能测试的简化版本
"""

import os
import sys
import argparse

def main():
    parser = argparse.ArgumentParser(description='PyG 算子性能测试')
    parser.add_argument('--device', type=str, default='auto', 
                       choices=['auto', 'cpu', 'cuda', 'npu'],
                       help='测试设备 (默认: auto)')
    parser.add_argument('--quick', action='store_true',
                       help='快速测试模式 (较少的测试配置)')
    parser.add_argument('--warmup', type=int, default=3,
                       help='热身迭代次数 (默认: 3)')
    parser.add_argument('--iterations', type=int, default=5,
                       help='测试迭代次数 (默认: 5)')
    parser.add_argument('--output', type=str, default=None,
                       help='输出文件名 (默认: 自动生成)')
    
    args = parser.parse_args()
    
    # 导入主要的性能测试模块
    try:
        from performance_test import PerformanceTester
        import torch
        import torch_npu
    except ImportError as e:
        print(f"导入错误: {e}")
        print("请确保已安装所有依赖项。")
        sys.exit(1)
    
    # 设备选择
    if args.device == 'auto':
        if hasattr(torch, 'npu') and torch.npu.is_available():
            device = 'npu'
        elif torch.cuda.is_available():
            device = 'cuda'
        else:
            device = 'cpu'
    else:
        device = args.device
    
    print(f"选择的设备: {device}")
    
    # 测试配置
    if args.quick:
        # 快速测试配置
        test_configs = [
            {'num_nodes': 100, 'num_edges': 500, 'node_features': 32, 'edge_features': 16},
            {'num_nodes': 1000, 'num_edges': 5000, 'node_features': 64, 'edge_features': 32},
            {'num_nodes': 5000, 'num_edges': 25000, 'node_features': 128, 'edge_features': 64},
        ]
        print("运行快速测试模式 (3个配置)")
    else:
        # 完整测试配置
        test_configs = [
            # 小规模图
            {'num_nodes': 100, 'num_edges': 500, 'node_features': 32, 'edge_features': 16},
            {'num_nodes': 500, 'num_edges': 2000, 'node_features': 64, 'edge_features': 32},
            
            # 中等规模图
            {'num_nodes': 1000, 'num_edges': 5000, 'node_features': 64, 'edge_features': 32},
            {'num_nodes': 2000, 'num_edges': 10000, 'node_features': 128, 'edge_features': 64},
            
            # 大规模图
            {'num_nodes': 5000, 'num_edges': 25000, 'node_features': 128, 'edge_features': 64},
            {'num_nodes': 10000, 'num_edges': 50000, 'node_features': 256, 'edge_features': 128},
        ]
        print("运行完整测试模式 (6个配置)")
    
    # 输出文件名
    if args.output is None:
        mode = "quick" if args.quick else "full"
        output_file = f"pyg_ops_performance_{device}_{mode}.json"
    else:
        output_file = args.output
    
    # 创建测试器并运行
    tester = PerformanceTester(
        device=device, 
        warmup_iterations=args.warmup, 
        test_iterations=args.iterations
    )
    
    # 运行测试
    results = tester.run_comprehensive_test(test_configs)
    
    # 生成报告
    tester.generate_report(results, output_file)
    
    print(f"\n性能测试完成！结果已保存到: {output_file}")

if __name__ == "__main__":
    main() 